import os
import numpy as np
import os.path as osp
import argparse
from PIL import Image
from scipy.io import loadmat


def mkdir_if_missing(directory):
    if not osp.exists(directory):
        os.makedirs(directory)


def extract_and_save(data, label, save_dir):
    for i, (x, y) in enumerate(zip(data, label)):
        if x.shape[2] == 1:
            x = np.repeat(x, 3, axis=2)
        if y == 10:
            y = 0
        x = Image.fromarray(x, mode='RGB')
        save_path = osp.join(
            save_dir,
            str(i + 1).zfill(6) + '_' + str(y) + '.jpg'
        )
        x.save(save_path)


def load_mnist(data_dir, raw_data_dir):
    filepath = osp.join(raw_data_dir, 'mnist_data.mat')
    data = loadmat(filepath)

    train_data = np.reshape(data['train_32'], (55000, 32, 32, 1))
    test_data = np.reshape(data['test_32'], (10000, 32, 32, 1))

    train_label = np.nonzero(data['label_train'])[1]
    test_label = np.nonzero(data['label_test'])[1]

    return train_data, test_data, train_label, test_label


def load_mnist_m(data_dir, raw_data_dir):
    filepath = osp.join(raw_data_dir, 'mnistm_with_label.mat')
    data = loadmat(filepath)

    train_data = data['train']
    test_data = data['test']

    train_label = np.nonzero(data['label_train'])[1]
    test_label = np.nonzero(data['label_test'])[1]

    return train_data, test_data, train_label, test_label


def load_svhn(data_dir, raw_data_dir):
    train = loadmat(osp.join(raw_data_dir, 'svhn_train_32x32.mat'))
    train_data = train['X'].transpose(3, 0, 1, 2)
    train_label = train['y'][:, 0]

    test = loadmat(osp.join(raw_data_dir, 'svhn_test_32x32.mat'))
    test_data = test['X'].transpose(3, 0, 1, 2)
    test_label = test['y'][:, 0]

    return train_data, test_data, train_label, test_label


def load_syn(data_dir, raw_data_dir):
    filepath = osp.join(raw_data_dir, 'syn_number.mat')
    data = loadmat(filepath)

    train_data = data['train_data']
    test_data = data['test_data']

    train_label = data['train_label'][:, 0]
    test_label = data['test_label'][:, 0]

    return train_data, test_data, train_label, test_label


def load_usps(data_dir, raw_data_dir):
    filepath = osp.join(raw_data_dir, 'usps_28x28.mat')
    data = loadmat(filepath)['dataset']

    train_data = data[0][0].transpose(0, 2, 3, 1)
    test_data = data[1][0].transpose(0, 2, 3, 1)

    train_data *= 255
    test_data *= 255

    train_data = train_data.astype(np.uint8)
    test_data = test_data.astype(np.uint8)

    train_label = data[0][1][:, 0]
    test_label = data[1][1][:, 0]

    return train_data, test_data, train_label, test_label


def main(data_dir):
    data_dir = osp.abspath(osp.expanduser(data_dir))
    raw_data_dir = osp.join(data_dir, 'Digit-Five')

    if not osp.exists(data_dir):
        raise FileNotFoundError('"{}" does not exist'.format(data_dir))

    datasets = ['mnist', 'mnist_m', 'svhn', 'syn', 'usps']

    for name in datasets:
        print('Creating {}'.format(name))

        output = eval('load_' + name)(data_dir, raw_data_dir)
        train_data, test_data, train_label, test_label = output

        print('# train: {}'.format(train_data.shape[0]))
        print('# test: {}'.format(test_data.shape[0]))

        train_dir = osp.join(data_dir, name, 'train_images')
        mkdir_if_missing(train_dir)
        test_dir = osp.join(data_dir, name, 'test_images')
        mkdir_if_missing(test_dir)

        extract_and_save(train_data, train_label, train_dir)
        extract_and_save(test_data, test_label, test_dir)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        'data_dir', type=str, help='directory containing Digit-Five/'
    )
    args = parser.parse_args()
    main(args.data_dir)
